import ml_collections

import reflow.configs.default as default_config

# ! 必须定义 get_config 函数


def get_config():
    config = default_config.get_config()

    # diffusers
    config.diffusers = diffusers = ml_collections.ConfigDict()
    diffusers.pipeline_ckpt = 'checkpoints/SD-1-4'
    diffusers.pipeline = 'stable_diffusion'
    diffusers.scheduler = 'dpm_solver_multi'
    diffusers.gradient_checkpointing = True
    diffusers.use_xformers = True
    diffusers.reflow_ckpt_inference = "" # 将推理的 pipeline unet 换成 reflow model
    diffusers.reflow_ckpt_train = "" # 训练用的 reflow model 从这个 ckpt 加载

    diffusers.num_inference_steps = 25
    diffusers.guidance_scale = 7.5
    time_interval = (0., 1.0)
    start_step, stop_step = (
        int(diffusers.num_inference_steps*time_interval[0]),
        int(diffusers.num_inference_steps*time_interval[1])
    )
    diffusers.start_step = start_step
    diffusers.stop_step = stop_step
    
    diffusers.use_redit = False
    diffusers.redit_pipeline = 'alt_diffusion'
    diffusers.redit_scheduler = 'dpm_solver_multi'
    diffusers.redit_pipeline_ckpt = 'checkpoints/AltDiffusion'
    diffusers.redit_inference_steps = 25
    diffusers.redit_guidance_scale = 7.5
    diffusers.redit_strength = 0.4

    # ema
    config.ema = ema = ml_collections.ConfigDict()
    ema.decay = 0.999

    # training
    training = config.training
    training.num_steps = 10
    training.batch_size = 1
    training.gradient_accumulation_steps = 1
    training.ckpt_path = ""  # NOTE only used when resume training
    training.p_uncond = 0.

    training.mixed_precision = 'fp16'
    training.reduce_mean = True

    training.log_freq = 10
    # training.eval_freq = 10
    # training.sampling_freq = 100
    # training.snapshot_freq = 1000

    # reflow
    config.reflow = reflow = ml_collections.ConfigDict()
    # NOTE: t0, t1, uniform, or an integer k > 1
    reflow.reflow_t_schedule = "uniform"
    reflow.reflow_loss = 'l2'  # NOTE: l2, lpips, lpips+l2, msssim+l1, l1, lpips+l1
    reflow.zt_compress_rate = 1.0
    
    reflow.finetune_sd = 'no'
    reflow.sd_t_schedule = ''

    # data
    data = config.data
    data.caption_path = 'data/reflow/laion6+_random1M.txt'
    data.dl_workers = 1
    # data.centered = True

    optim = config.optim
    optim.use_8bit_adam = False
    optim.lr_scheduler = 'constant'
    optim.lr = 1e-5
    optim.warmup = 0

    config.seed = 718694
    # config.device=None

    return config
